package com.origin.sse.service.impl;

import cn.hutool.core.text.CharSequenceUtil;
import com.origin.sse.constant.EventConst;
import com.origin.sse.service.SseEmitterService;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * Sse(server-send event)
 *
 * @author yezh
 * @date 2024/1/15 19:23
 */
@Service
@Slf4j
public class SseEmitterServiceImpl implements SseEmitterService {

    @Value("${sse.timeout:5}")
    private long sseTimeout;

    private static final Map<String, SseEmitterWrapper> sseCache = new ConcurrentHashMap<>();


    @Override
    public SseEmitter getSseConn(String clientId, List<String> eventTypeList) {
         SseEmitter sseEmitter = sseCache
                .computeIfAbsent(clientId, k -> new SseEmitterWrapper(clientId, createSseEmitter(clientId), new HashSet<>(eventTypeList)))
                .getSseEmitter();
         log.info("当前用户数: {}", sseCache.size());
         return sseEmitter;
    }

    @Override
    public void sendMsgToClient(String data, String eventType) {
        if (CharSequenceUtil.isNotBlank(data) && CharSequenceUtil.isNotBlank(eventType)) {
            sseCache.forEach((clientId, sseEmitterWrapper) -> {
                if (sseEmitterWrapper.containsDataType(eventType)) {
                    sendMsgToClientByClientId(clientId, data, eventType, sseEmitterWrapper.getSseEmitter());
                }
            });
        }
    }

    @Override
    public void sendMsgToClientByClientId(String clientId, String data, String eventType, SseEmitter sseEmitter) {
        if (sseEmitter == null) {
            log.error("SseEmitterServiceImpl[sendMsgToClient]: 推送消息失败：客户端 {} 未创建长链接,失败消息: {}",
                    clientId, data);
            return;
        }

        SseEmitter.SseEventBuilder sendData = SseEmitter.event().name(eventType).data(data, MediaType.APPLICATION_JSON);
        try {
            sseEmitter.send(sendData);
        } catch (IOException e) {
            // 推送消息失败，记录错误日志，进行重推
            log.error("SseEmitterServiceImpl[sendMsgToClient]: 推送消息失败：{},尝试进行重推", data, e);
            boolean isSuccess = true;
            // 推送消息失败后，每隔10s推送一次，推送5次
            for (int i = 0; i < 5; i++) {
                try {
                    Thread.sleep(10000);
                    SseEmitterWrapper wrapper = sseCache.get(clientId);
                    if (wrapper != null) {
                        sseEmitter = wrapper.getSseEmitter();
                        if (sseEmitter == null) {
                            log.error("SseEmitterServiceImpl[sendMsgToClient]：{}的第{}次消息重推失败，未创建长链接", clientId, i + 1);
                            continue;
                        }
                        sseEmitter.send(sendData);
                    }
                } catch (Exception ex) {
                    log.error("SseEmitterServiceImpl[sendMsgToClient]：{}的第{}次消息重推失败", clientId, i + 1, ex);
                    continue;
                }
                log.info("SseEmitterServiceImpl[sendMsgToClient]：{}的第{}次消息重推成功,{}", clientId, i + 1, data);
                return;
            }
        }
    }

    @Override
    public void closeSseConnect(String clientId) {
        if (sseCache.containsKey(clientId)) {
            SseEmitterWrapper sseEmitterWrapper = sseCache.remove(clientId);
            sseEmitterWrapper.getSseEmitter().complete();
        }
    }

    /**
     * 长链接完成后回调接口(即关闭连接时调用)
     */
    private Runnable completionCallBack(String clientId) {
        return () -> {
            log.info("结束连接：{}", clientId);
            removeUser(clientId);
        };
    }

    private SseEmitter createSseEmitter(String clientId) {
        // 设置超时时间, 单位: mm, 0 表示不过期. 默认 30 秒, 超过时间未完成会抛出异常：AsyncRequestTimeoutException
        SseEmitter sseEmitter = new SseEmitter(sseTimeout * 1000L);
        // 注册回调
        sseEmitter.onCompletion(completionCallBack(clientId));
        sseEmitter.onTimeout(timeoutCallBack(clientId));
        sseEmitter.onError(errorCallBack(clientId));
        log.info("创建新的 sse 连接，当前用户：{}", clientId);

        try {
            sseEmitter.send(SseEmitter.event().name(EventConst.CLIENT_ID).data(clientId));
        } catch (IOException e) {
            log.error("SseEmitterServiceImpl[createSseConnect]: 创建长链接异常，客户端 ID: {}", clientId, e);
            throw new RuntimeException("创建连接异常！", e);
        }
        return sseEmitter;
    }

    /**
     * 连接超时
     */
    private Runnable timeoutCallBack(String clientId) {
        return () -> {
            log.info("连接超时：{}", clientId);
            removeUser(clientId);
        };
    }

    /**
     * 推送异常时
     */
    private Consumer<Throwable> errorCallBack(String clientId) {
        return throwable -> {
            log.error("SseEmitterServiceImpl[errorCallBack]：连接异常,客户端ID:{}", clientId);

            // 推送消息失败后，每隔 10s 推送一次, 推送 5 次
            for (int i = 0; i < 5; i++) {
                try {
                    TimeUnit.SECONDS.sleep(10);
                    SseEmitterWrapper wrapper = sseCache.get(clientId);
                    if (wrapper == null || wrapper.getSseEmitter() == null) {
                        log.error("SseEmitterServiceImpl[errorCallBack]：第{}次消息重推失败,未获取到 {} 对应的长链接", i + 1, clientId);
                        continue;
                    }
                    wrapper.getSseEmitter().send("失败后重新推送");
                } catch (InterruptedException e) {
                    log.error("中断异常", e);
                } catch (IOException e) {
                    log.error("IO 异常", e);
                }
            }
        };
    }

    private void removeUser(String clientId) {
        sseCache.remove(clientId);
        log.info("SseEmitterServiceImpl[removeUser]:移除用户：{}", clientId);
    }

    @Data
    @AllArgsConstructor
    static class SseEmitterWrapper {
        private String clientId;
        private SseEmitter sseEmitter;
        private Set<String> dataTypeList;

        public boolean containsDataType(String dataType) {
            return dataTypeList.contains(dataType);
        }

    }

}
